import numpy as np
import scipy.special as sc_special

# 本次修改将参数前置
nest_num = 50
nest_dim = 2
iter_num = 10000
pa = 0.25
beta = 1.5  # 莱维飞行公式的β
step_coefficient = 0.01  # 步长控制因子

version = '1.1.0'


# 生成鸟窝
def generate_neste(nest_num, nest_dim, lower_boundary, upper_boundary):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    nests = np.empty((nest_num, nest_dim))

    # 遍历生成鸟窝，此问题相当于生成x，y
    for each_nest in range(nest_num):
        # 下边界 + 系数 * (上-下)
        nests[each_nest] = lower_boundary + np.array([np.random.rand() for _ in range(nest_dim)]) * (
                upper_boundary - lower_boundary)

    return nests


# 适应度计算函数
def calc_fitness(fit_func, nests):
    finess = np.empty(nest_num)

    # 遍历nests 计算 适应度 填入 finess
    for each_nest in range(nest_num):
        finess[each_nest] = fit_func(nests[each_nest])

    return finess


# 莱维飞行函数,完全按照论文复现
def levy_flight(nest_num, nest_dim, beta):
    sigma_u = (sc_special.gamma(1 + beta) * np.sin(np.pi * beta / 2) / (
            sc_special.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2)))) ** (1 / beta)
    sigma_v = 1

    # 正态分布参数：均值，标准差，shape
    u = np.random.normal(0, sigma_u, (nest_num, nest_dim))
    v = np.random.normal(0, sigma_v, (nest_num, nest_dim))

    steps = u / ((np.abs(v)) ** (1 / beta))

    return steps


# 更新鸟窝函数
def update_nests(fit_func, lower_boundary, upper_boundary, nests, best_nest, fitness, step_coefficient):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    steps = levy_flight(nest_num, nest_dim, beta)

    new_nests = nests.copy()

    for each_nest in range(nest_num):
        # 步长的计算，步长系数*步长*(鸟窝-最佳鸟窝)
        step_size = step_coefficient * steps[each_nest] * (nests[each_nest] - best_nest)

        # 随机生成方向
        # step_direction = np.random.rand(nest_dim)
        # 更新nest,原始步长+步长*方向
        # new_nests[each_nest] += step_size * step_direction

        # 没有随机方向的
        new_nests[each_nest] += step_size
        # 新生成的nest小于下界即取下界
        new_nests[each_nest][new_nests[each_nest] < lower_boundary] = lower_boundary[
            new_nests[each_nest] < lower_boundary]
        # 新生成的nest大于上界即取上界
        new_nests[each_nest][new_nests[each_nest] > upper_boundary] = upper_boundary[
            new_nests[each_nest] > upper_boundary]

    # 求新的nest的适应值
    new_fitness = calc_fitness(fit_func, new_nests)
    # 新的适应值大于 旧适应值 取新适应值 对应的nest
    nests[new_fitness > fitness] = new_nests[new_fitness > fitness]

    return nests


# 遗弃鸟窝的函数，也就是二次更新函数，没有加论文中的跃进函数
def abandon_nests(nests, lower_boundary, upper_boundary, pa):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    for each_nest in range(nest_num):
        # 生成随机数来控制遗弃的概率,生成随机数的范围为0-1
        if (np.random.rand() < pa):
            # 步长 = 随机数* (俩个随机nest的差)，除了加跃进函数，完全严格按照论文复现
            step_size = np.random.rand() * (
                    nests[np.random.randint(0, nest_num)] - nests[np.random.randint(0, nest_num)])
            nests[each_nest] += step_size
            # 越界取值判断，同理，超过上下界取上下界
            nests[each_nest][nests[each_nest] < lower_boundary] = lower_boundary[nests[each_nest] < lower_boundary]
            nests[each_nest][nests[each_nest] > upper_boundary] = upper_boundary[nests[each_nest] > upper_boundary]

        return nests


def run(nest_num, nest_dim, fit_func, lower_boundary, upper_boundary, iter_num, pa, beta,
        step_coefficient):
    # 生成鸟窝nest-x,y
    nests = generate_neste(nest_num, nest_dim, lower_boundary, upper_boundary)

    # 计算适应度
    fitness = calc_fitness(fit_func, nests)

    # 记录最佳适应度的位置
    # 论文为最小优化
    best_nest_index = np.argmin(fitness)

    # 记录最佳适应度的值
    best_fitness = fitness[best_nest_index]
    # 记录最佳适应度对应的nest的值(鸟窝)
    best_nest = nests[best_nest_index].copy()

    for _ in range(iter_num):
        # 更新鸟窝
        nests = update_nests(fit_func, lower_boundary, upper_boundary, nests, best_nest, fitness, step_coefficient)
        # 遗弃鸟窝，本质就是二次跟新鸟窝nest，也就是二次更新x,y
        nests = abandon_nests(nests, lower_boundary, upper_boundary, pa)
        # 计算适应度
        fitness = calc_fitness(fit_func, nests)


        # 根据论设置为为最小值
        min_nest_index = np.argmin(fitness)
        # 最小适应度的值
        min_fitness = fitness[min_nest_index]
        # 最小适应度对应的nest -- x，y
        min_nest = nests[min_nest_index]

        # 判断，如果该次的最大适应度大于之前的，就换掉
        if (min_fitness < best_fitness):
            best_nest = min_nest.copy()
            best_fitness = min_fitness
        # 返回最佳nest--xy，和最大适应度
    return (best_nest, best_fitness)


if __name__ == '__main__':
    def fit_func(nest):
        x, y = nest
        return 200 - (x ** 2 + y - 11) ** 2 - (x + y ** 2 - 7) ** 2


    result = []
    # for i in range(10):
    best_nest, best_fitness = run(nest_num, nest_dim, fit_func, [-3, -3], [3, 3], iter_num, pa, beta,
                                  step_coefficient)

    # print('最大值为:%.5f, 在(%.5f, %.5f)处取到!' % (best_fitness, best_nest[0], best_nest[1]))
    print('最小值为:%.5f, 在(%.5f, %.5f)处取到!' % (best_fitness, best_nest[0], best_nest[1]))
#     print(result)
